import pickle, gzip
import sys, inspect

###Exception classes###
class InvalidArgError(Exception):
    def __init__(self, value):
        self.parameter = value
    def __str__(self):
        return repr(self.parameter)

class VarDoesNotExist(Exception):
    def __init__(self, var_name):
        self.var_name = var_name
    def __str__(self):
        return 'error: Argument [%s] not found in calling frame'%self.var_name

###MAIN FUNCTIONS###
def save(fname, *args):
    """
    description:
        Saves variables specified in args to a file specified by fname.
        The file is also compressed using the gzip library.
    input:
        fname (string) - path to file to save data to.
        *args (string(s)) - name(s) of variables to save
    output:
        None
    side effects:
        creates file specified by fname. 
    example usage:
        x = 1
        y = 2
        save('temp.dat', 'x', 'y')
    """
    #check to see if fname is a string
    if not(type(fname) is str):
        raise InvalidArgError('error: fname must be a string')
    #get caller's local variables
    caller_lv = get_caller_locals()
    #collect all variables specified by args into a dictionary
    var_dict = dict()
    for argu in args:    
        #error check -- check if argument is a string
        if not(type(argu) is str):
            raise InvalidArgError('error: Argument must be of type str')
        #error check -- check if argument is in caller's stackframe (local vars)
        argu = argu.strip() #strip off excess whitespace
        if not(argu in caller_lv):
            raise VarDoesNotExist(argu)
        #argument is valid so add its data to dictionary
        var_dict[argu] = caller_lv[argu]
    #pickle dictionary to file
    target_file = gzip.open(fname, 'w')
    pickle.dump(var_dict, target_file)
    target_file.close()

def load(fname):
    """
    description:
        Load data stored in fname into the caller's stack frame.
    input:
        fname (string) - path to compressed data file.
    output:
        dictionary of variables saved
    side effects:
        modifies caller's stack frame to include variables epcified in the
        data file. NOTE: this only works if load is called from the top most 
                          frame. It currently can NOT modify the local
                          variables if called from within a function or method.
    example usage:
        x = 1
        y = 2
        save('temp.dat', 'x', 'y')
        x = None
        Y = None
        load('temp.dat')
        print x
        print y    
    """
    #check to see if fname is a string
    if not(type(fname) is str):
        raise InvalidArgError('error: fname must be a string')

    #get caller's local variables
    caller_lv = get_caller_locals()
    #load dictionary from fname 
    target_file = gzip.open(fname,'r')
    pdict = pickle.load(target_file)
    #update caller's local variables with values in pickled dictionary
    caller_lv.update(pdict)
    target_file.close()
    return pdict

def get_caller_locals():
    """
    description:
        get the local variables of the caller of the function which 
        invokes get_caller_locals()
    input:
        None
    output:
        local variable dictionary of the caller's caller's stack frame 
        (e.g., the variables of the caller of the function which envoked 
        get_caller_locals())
    side effects:
        None
    example usage:
        def foo():
            d = get_caller_locals()
            assert(d['x'] == 100)
            print d['x']
        x = 100
        foo() 
        x = 99
        foo()   #assertion should raise exception here
    """
    #caller_lv =  sys._current_frames().values()[-1].f_back.f_back.f_locals
    caller_lv = inspect.currentframe().f_back.f_back.f_locals
    return caller_lv
